# -*- coding: utf-8 -*-
# train.py

# from scipy.sparse import csr_matrix

from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB, GaussianNB, BernoulliNB
from sklearn.model_selection import train_test_split
from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import GridSearchCV
# from sklearn.feature_extraction.text import TfidfTransformer
import pickle
import sys
import numpy as np


if __name__ == '__main__':
    data_file = sys.argv[1]
    model_file = sys.argv[2]

    X, target = pickle.load(open(data_file, 'rb'))

    train_x, test_x, train_y, test_y = train_test_split(X, target, test_size=0.20, random_state=30)
    nb = BernoulliNB()

    nb.fit(train_x, train_y)

    # 保存训练好的模型
    with open(model_file, 'wb+') as fd:
        pickle.dump(nb, fd)

    for k, v in zip(nb.predict(test_x[0:100]), test_y[0:100]):
        print("%s %s" % (k,v))

    print(nb)

    err_cnt = sum(np.array(nb.predict(test_x)) != np.array(test_y))
    print(err_cnt)
    print(err_cnt/len(test_y))

